#!/usr/bin/env python
import atlasgen as atl
import cPickle
import os
from mpl_toolkits.mplot3d import Axes3D
from numpy import linalg as LA
from itk import AffineTransform as at
from sys import stdin as inp
import matplotlib as mpl
import matplotlib.pyplot as pp
import numpy as np
import imageTools.ITKUtils.io as io
import itk
import scipy.ndimage as ndi
import math
import imageTools.browse.a3d as a3d
import sys
import itertools as it
import random

def main():
	

	gfo = file(sys.argv[1],'rb')
	graph = cPickle.load(gfo)

	pulm = file(sys.argv[2],'rb')
	pulmatlas = cPickle.load(pulm)

	aorta = file(sys.argv[3],'rb')
	aortaatlas = cPickle.load(aorta)
	
	test = file(sys.argv[4],'w')
	test.write("Graph tested: {0}\n".format(sys.argv[1]) )
	test.write("pulmonary atlas: {0} , aorta atlas: {1} \n".format(sys.argv[2],sys.argv[3]) )

	trials = 20
	ispulm =0
	isaorta =0
	pulmscores = []
	aortscores = []
	bestpulm = float(0)
	bestaort=float(0)
	margins = []

	for trial in range(trials):	
		
		#optimize rotations/membership
		pulm_rotation = atl.autorotate(graph,pulmatlas,iterate=60)
		pulm_membership = atl.score_by_rank(pulm_rotation,pulmatlas)

		aorta_rotation = atl.autorotate(graph,pulmatlas,iterate=60)
		aorta_membership = atl.score_by_rank(aorta_rotation,aortaatlas)

	
		#print "Pulm membership:",pulm_membership," Aorta membership:", aorta_membership

		#pulm_membership = atl.graph_score(pulm_rotation,pulmatlas,atl.closest_node)
		#aorta_membership = atl.graph_score(aorta_rotation,aortaatlas,atl.closest_node)	
	
		#keep score and update lists
		membership = max(pulm_membership,aorta_membership)
		pulmscores.append(pulm_membership)
		aortscores.append(aorta_membership)
		
		if pulm_membership > bestpulm:
			pgraph = pulm_rotation
			bestpulm = pulm_membership

		if aorta_membership > bestaort:
			agraph = aorta_rotation
			bestaort = aorta_membership

		#comparison
		print "Pulm membership:",pulm_membership," Aorta membership:", aorta_membership
		test.write("Trial: {0}, Pulm membership: {1}, Aorta membership: {2} \n".format(trial,pulm_membership,aorta_membership) )

		if membership == pulm_membership:
			print "is pulmonary"
			test.write("Is pulmonary, pulm count: {0}\n".format(ispulm) )
			pmean = np.mean(pulmscores)
			pmargin = np.std(pulmscores)
			ispulm = ispulm+1
			print "pulm count:",ispulm,"avg membership=",pmean, "+/-",pmargin
		else:	
			print "is aorta"
			test.write("Is aorta, aorta count: \n".format(isaorta) )
			amean = np.mean(aortscores)
			amargin = np.std(aortscores)
			isaorta = isaorta+1
			print "aorta count:",isaorta,"avg membership",amean,"+/-:",amargin
	
	#statistics
	print "Avg pulm membership:",pmean,"Avg aorta membership:", amean
	test.write("Avg pulm membership:{0}, Avg aorta membership: {1} \n".format(pmean,amean) )

	#calculate probability/percent error (assume margins are uniformly distributed)

	lowerbound = min((pmean-pmargin),(amean-amargin))
	upperbound = max((pmean+pmargin),(amean+amargin))

	erupper = min((pmean+pmargin),(amean+amargin))
	erlower = max((pmean-pmargin),(amean-amargin))
	
	error = erupper-erlower
	if error < 0:	error = 0	#then there is no overlap
	else:	pass

	rangebounds = upperbound-lowerbound
	perror = error/rangebounds

	print "Probability of error =",perror
	test.write("Probability of error = {0}\n".format(perror))

	#decision, with certainty
	
	if ispulm>isaorta:
		certainty= float(ispulm)/float(trials)
		atlas = pulmatlas
		rotatedgraph = pgraph
		print "is pulmonary with certainty:",certainty
		test.write("Is pulm with certainty {0}".format(certainty) )
	else:
		certainty = float(isaorta)/float(trials)
		atlas = aortaatlas
		rotatedgraph = agraph
		print "is aorta with certainty:",certainty
		test.write("Is aorta with certainty {0}".format(certainty) )

	test.close()

	window = [(-30,70), (-30,70), (-30,70)]	
	scale = 6
	
	atl.plotatlasandgraph(atlas,rotatedgraph,window,scale)
	
	return 0

if __name__== "__main__":
	main()
#-----------------------END OF PROGRAM------------------------
